import datetime
import time

import numpy as np
import cv2

global_intrusion_record = []

# Add the recognition you want to detect
global_active_objects = ['person']

def object_detection():
    global global_intrusion_record

    # Set colors
    colors = []
    for i in range(91):
        random_c = np.random.randint(256, size=3)
        colors.append((int(random_c[0]), int(random_c[1]), int(random_c[2])))

    # Opencv DNN
    net = cv2.dnn.readNet("model/object_detect/yolov4-tiny.weights", "model/object_detect/yolov4-tiny.cfg")
    model = cv2.dnn_DetectionModel(net)
    model.setInputParams(size=(320, 320), scale=1 / 255)

    # Load class lists
    classes = []
    with open("model/object_detect/classes.txt", "r") as file_object:
        for class_name in file_object.readlines():
            class_name = class_name.strip()
            classes.append(class_name)



    # Local camera /  fps: 25
    cap = cv2.VideoCapture(0)
    # print(cap.get(5))

    cnt = 0
    person_num = 0
    while True:
        # 抽帧处理
        cnt = cnt + 1
        ret = cap.grab()
        if ret is False:
            continue

        _, frame = cap.retrieve()

        # Object Detection
        (class_ids, scores, bboxes) = model.detect(frame, confThreshold=0.3, nmsThreshold=.4)
        for class_id, score, bbox in zip(class_ids, scores, bboxes):
            (x, y, w, h) = bbox
            class_name = classes[class_id]
            color = colors[class_id]
            if class_name in global_active_objects:
                if cnt % 200 == 0:
                    person_num = person_num + 1
                cv2.putText(frame, class_name, (x, y - 10), cv2.FONT_HERSHEY_PLAIN, 2, color, 2)
                cv2.rectangle(frame, (x, y), (x + w, y + h), color, 5)
        # 每200帧就统计一下当前画面检测到的人数，并把画面保存成jpg
        if cnt % 200 == 0:
            t = str(datetime.datetime.now().replace(microsecond=0)).replace(' ', 'T').replace(':', '-')
            ret, jpeg = cv2.imencode('.jpg', frame)
            if ret:
                path = 'resource/intrusion_image/' + str(t) + '.jpg'
                with open(path, 'wb') as f:
                    f.write(jpeg.tobytes())
                record = {'cnt': cnt, 'time': t, 'number': person_num, 'path': path}
                global_intrusion_record.append(record)
            person_num = 0
            yield frame, cnt
        else:
            yield frame, None

def get_intrusion_record():
    global global_intrusion_record
    return global_intrusion_record

def set_active_objects(new_active_objects):
    global global_active_objects
    global_active_objects = new_active_objects
